import config
import pandas as pd
import matplotlib.pyplot as plt

id2label = {}
with open(config.label_ids_file, 'r', encoding='utf-8') as txt:
    for line in txt:
        ID, label = line.strip().split('\t')
        id2label[ID] = label

# print(id2label)

for filePath in [config.train_raw_file, config.eval_raw_file, config.test_raw_file]:
    samples = []
    with open(filePath, 'r', encoding='utf-8') as txt:
        for line in txt:
            ID, text = line.strip().split('\t')
            label = id2label[ID]
            sample = label + '\t' + text
            samples.append(sample)

        outfile = config.train_data_file
        if 'eval' in filePath:
            outfile = config.eval_data_file
        if 'test' in filePath:
            outfile = config.test_data_file
        with open(outfile, 'w', encoding='utf-8') as csv:
            csv.write('label\ttext\n')
            for sample in samples:
                csv.write(sample)
                csv.write('\n')
